# -*- coding: utf-8 -*-

# 导入pyspark
from pyspark import SparkContext
from pyspark.mllib.classification import NaiveBayes
from pyspark.mllib.regression import LabeledPoint
from pyspark.mllib.linalg import Vectors
from pyspark.mllib.tree import DecisionTree
import numpy as np

# 导入可视化
import matplotlib
matplotlib.use('Agg') # 不回显
import matplotlib.pyplot as plt

# 计时工具
import time

# 朴素贝叶斯分类模型
def naiveBayes(trainData, testData):
	# startTime = time.time()
	model = NaiveBayes.train(trainData, 1.0)
	predictionAndLabel = testData.map(lambda p: (model.predict(p.features), p.label))
	accuracy = 1.0 * predictionAndLabel.filter(lambda (x, v): x == v).count() / testData.count()
	# endTime = time.time()
	# durationTime = endTime - startTime
	# return durationTime, accuracy
	return accuracy

# 决策树分类模型
def decisionTree(trainData, testData, numClasses, categoricalFeaturesInfo, maxBins, maxDepth):
	# startTime = time.time()
	modelDecisionTree = DecisionTree.trainClassifier(trainData, numClasses, categoricalFeaturesInfo , maxDepth = maxDepth, maxBins = maxBins)
	predictDecisionTree = modelDecisionTree.predict(testData.map(lambda p: p.features))
	predictionAndLabelDecisionTree = testData.map(lambda p: p.label).zip(predictDecisionTree).map(lambda (x, y): (y, x))
	accuracyDecisionTree = 1.0 * predictionAndLabelDecisionTree.filter(lambda (x, v): x == v).count() / testData.count()
	# endTime = time.time()
	# durationTime = endTime - startTime
	# return durationTime, accuracyDecisionTree
	return accuracyDecisionTree

# 画图
def drawFigure(xArr, yArr, title):
	figName = "%s.png"%title
	plt.figure(num=1, figsize=(8,6))
	plt.title(title, size=14)
	plt.xlabel('x', size=14)
	plt.ylabel('y', size=14)
	plt.plot(xArr, yArr)
	plt.savefig(figName, format='png')

# main函数部分
sc = SparkContext("yarn-client", "NaiveBayes Spark App")

# 加载数据集 
dataPath = "hdfs://192.168.119.141:9100/data/classification/mat_classify"
# 原始数据集
raw_data = sc.textFile(dataPath)
# LabeledPoint数据集
records = raw_data.map(lambda x: x.split(","))
labeledRecords = records.map(lambda x: LabeledPoint(float(x[0]), Vectors.dense([float(y) for y in x[1].split(" ")])))

# 拆分训练集与测试集
data_with_idx = labeledRecords.zipWithIndex().map(lambda (k, v): (v, k))
testRecords = data_with_idx.sample(False, 0.2, 42)
trainRecords = data_with_idx.subtractByKey(testRecords)
testData = testRecords.map(lambda (idx, p): p)
trainData = trainRecords.map(lambda (idx, p): p)
trainData.cache()
testData.cache()

# 写入文件

# 训练朴素贝叶斯分类模型
naiveBayes(trainData, testData)

# 训练决策树分类模型
# 最大深度
depthParams = [1, 2, 3, 4, 5, 10, 20]
depthMetrics = [decisionTree(trainData, testData, 3, {}, 32, maxDepth) for maxDepth in depthParams]
drawFigure(depthParams, depthMetrics, 'DecisionTreeClassificationModelDepth')

# 最大划分数
binsParams = [2, 4, 8, 16, 32, 64]
binsMetrics = [decisionTree(trainData, testData, 3, {}, maxBins, 5) for maxBins in binsParams]
drawFigure(binsParams, binsMetrics, 'DecisionTreeClassificationModelBins')

# 关闭sc
sc.stop()